package com.lei.study.lock;

import org.apache.zookeeper.CreateMode;
import org.apache.zookeeper.KeeperException;
import org.apache.zookeeper.ZooDefs;
import org.apache.zookeeper.ZooKeeper;
import org.apache.zookeeper.data.Stat;

import java.io.IOException;
import java.util.Collections;
import java.util.List;
import java.util.concurrent.CountDownLatch;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicInteger;

/**
 * 使用zookeeper实现分布式锁
 *
 * @author 伍磊
 */
public class ZkLock implements Lock{

   private static final String ZK_LOCK_PREFIX = "/lock";

    /**
     * 加锁的模块
     */
   private String lockModule = "/test";

   private static final String SPE = "/";

   private static final String SEQ = "/seq-";

    /**
     * 重入锁
     */
   private AtomicInteger lockCount = new AtomicInteger(0);

    /**
     * 获得锁的线程
     */
   private Thread thread;

    /**
     * zookeeper连接客户端
     */
   private ZooKeeper zooKeeper;

    /**
     * 创建的节点
     */
   private String lockedPath;

    /**
     * 创建的节点排序号
     */
    private String lockedShortPath;

    /**
     * 临时节点的父节点
     */
    private String parentPath;

    /**
     * 前一个节点path
     */
    private String priorPath;

    /**
     * 加锁等待时间
     */
    private static final int WAIT_TIME = 3;

   public ZkLock() {
       init();
   }

    public ZkLock(String lockModule) {
        // 校验
        if (!lockModule.startsWith(SPE)) {
            System.err.println("格式错误!");
            throw new RuntimeException("格式错误!!");
        }
        this.lockModule = lockModule;
        init();
    }

    private void init() {
        // 1. 连接zookeeper
        // 使用原生zookeeper
        try {
            zooKeeper = new ZooKeeper("localhost:2181,localhost:2182,localhost:2183,localhost:2184",
                    1000,
                    event -> {
                        System.out.println("连接监听事件.... " + event.toString());
                    });
        } catch (IOException e) {
            System.err.println("初始化zookeeper连接异常: " + e.getMessage());
            e.printStackTrace();
            return;
        }

        // 2. 创建ZK_LOCK_PREFIX节点
        try {
            Stat exists = zooKeeper.exists(ZK_LOCK_PREFIX, false);
            if (null == exists) {
                // 创建root节点
                zooKeeper.create(ZK_LOCK_PREFIX, "lock".getBytes(), ZooDefs.Ids.OPEN_ACL_UNSAFE, CreateMode.PERSISTENT);
            }

            // 创建模块节点
            String node = ZK_LOCK_PREFIX + lockModule;
            Stat nodeExists = zooKeeper.exists(node, false);
            if (null == nodeExists) {
                // 创建root节点
                zooKeeper.create(node, "lock.module".getBytes(), ZooDefs.Ids.OPEN_ACL_UNSAFE, CreateMode.PERSISTENT);
            }

            parentPath = ZK_LOCK_PREFIX + lockModule;
            System.out.println("init success...");

        } catch (KeeperException | InterruptedException e) {
            System.err.println(e.getMessage());
            e.printStackTrace();
        }
    }


    @Override
    public boolean lock() {
        // 保证可重入锁
        synchronized (this) {
            if (lockCount.get() == 0) {
                thread = Thread.currentThread();
                lockCount.incrementAndGet();
            } else {
                // 若不是获取锁的线程来获取则返回false
                if (!thread.equals(Thread.currentThread())) {
                    return false;
                }

                int newCount = lockCount.incrementAndGet();
                System.out.println(thread.getName() + "可重入锁获取成功，count:" + newCount);
                return true;
            }
        }

        try {
            boolean locked;
            // 尝试加锁
            locked = tryLock();
            if (locked) {
                return true;
            }

            // 加锁失败则等待
            while (!locked) {
                // 等待 阻塞
                // 直到监听的前一个节点被删除 则走下面的代码判断当前节点是否为最小的
                await();

                // 获取等待的子节点列表
                List<String> waiters = getWaiters();

                // 判断是否加锁成功
                if (checkLocked(waiters)) {
                    locked = true;
                }
            }

            return true;
        } catch (Exception e) {
            e.printStackTrace();
            unLock();
        }

        return false;
    }

    @Override
    public boolean unLock() {
        // 只有加锁的线程可以解锁
        if (!thread.equals(Thread.currentThread())) {
            return false;
        }

        // 减少可重入的次数
        int newCount = lockCount.decrementAndGet();

        if (newCount < 0) {
            throw new IllegalMonitorStateException("计数错误: " + lockedPath);
        }

        // 计数器不为0 直接返回true
        if (newCount != 0) {
            System.out.println(thread.getName() + "释放可重入锁成功，count:" + newCount);
            return true;
        }

        // 为0 则需要删除临时节点
        try {
            zooKeeper.delete(lockedPath, -1);
        } catch (InterruptedException | KeeperException e) {
            e.printStackTrace();
            return false;
        }
        System.out.println(thread.getName() + "临时节点删除成功: " + lockedPath);
        return true;
    }

    /**
     * 监听前一个节点的删除
     */
    private void await() throws KeeperException, InterruptedException {
        if (null == priorPath) {
            throw new RuntimeException("priorPath error...");
        }

        final CountDownLatch countDownLatch = new CountDownLatch(1);

        zooKeeper.exists(priorPath, event -> {
            System.out.println("priorPath节点删除: " + priorPath);
            countDownLatch.countDown();
        });

        // 最长等待时间3秒
        countDownLatch.await(WAIT_TIME, TimeUnit.SECONDS);
    }

    /**
     * 尝试加锁
     * 未获取到则返回false
     * @return
     */
    private boolean tryLock() throws KeeperException, InterruptedException {
        // 获取加锁路径
        String lockPath = ZK_LOCK_PREFIX + lockModule + SEQ;

        // 创建临时节点
        lockedPath = zooKeeper.create(lockPath, "lock-se".getBytes(), ZooDefs.Ids.OPEN_ACL_UNSAFE, CreateMode.EPHEMERAL_SEQUENTIAL);

        if (null == lockedPath || "".endsWith(lockedPath)) {
            throw new RuntimeException("zk error!");
        }

        // 获取加锁的排队编号
        lockedShortPath = getShortName(lockedPath);

        // 获取所有等待的子节点列表
        List<String> waiters = getWaiters();

        // 判断自己是否为第一个
        if (checkLocked(waiters)) {
            return true;
        }

        // 未获取到锁 判断当前index
        int index = Collections.binarySearch(waiters, lockedShortPath);

        if (index < 0) {
            // 网络抖动 获取到的子节点中没有自己
            throw new RuntimeException("节点没有找到: " + lockedShortPath);
        }

        // 保存前一个节点的名称
        priorPath = parentPath + waiters.get(index - 1);

        return false;
    }

    /**
     * 判断自己是否为等待子节点中最小的节点 若是则获取锁成功
     * @param waiters
     * @return
     */
    private boolean checkLocked(List<String> waiters) {
        // 排序
        Collections.sort(waiters);

        String s = waiters.get(0);

        if (lockedShortPath.equals(s)) {
            System.out.println(thread.getName() + "成功获取到了锁...节点为: " + lockedShortPath);
            return true;
        }
        return false;
    }

    /**
     * 获取所有等待的子节点
     * @return
     */
    private List<String> getWaiters() throws KeeperException, InterruptedException {
        return zooKeeper.getChildren(parentPath, false);
    }

    /**
     * 获取节点short名称
     * @param lockedPath
     * @return
     */
    private String getShortName(String lockedPath) {
        int i = lockedPath.lastIndexOf('/');
        return lockedPath.substring(i + 1);
    }
}
